package com.pestdetection;

import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.conf.layers.SubsamplingLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

/**
 * 病虫害检测模型类
 * 使用卷积神经网络(CNN)来识别病虫害
 * 主要功能：
 * 1. 构建CNN模型
 * 2. 训练模型
 * 3. 获取训练好的模型
 */
public class PestDetectionModel {
    private MultiLayerNetwork 模型;  // 神经网络模型
    private final int 病虫害种类数;  // 需要识别的病虫害种类数量
    private final int 图片高度;      // 输入图片的高度
    private final int 图片宽度;      // 输入图片的宽度
    private final int 图片通道数;    // 输入图片的通道数（RGB为3）

    /**
     * 构造函数
     * @param 病虫害种类数 需要识别的病虫害种类数量
     * @param 图片高度 输入图片的高度
     * @param 图片宽度 输入图片的宽度
     * @param 图片通道数 输入图片的通道数
     */
    public PestDetectionModel(int 病虫害种类数, int 图片高度, int 图片宽度, int 图片通道数) {
        this.病虫害种类数 = 病虫害种类数;
        this.图片高度 = 图片高度;
        this.图片宽度 = 图片宽度;
        this.图片通道数 = 图片通道数;
        构建模型();
    }

    /**
     * 构建CNN模型
     * 模型结构：
     * 1. 卷积层：提取图片特征
     * 2. 池化层：降低特征维度
     * 3. 全连接层：整合特征
     * 4. 输出层：分类预测
     */
    private void 构建模型() {
        MultiLayerConfiguration 配置 = new NeuralNetConfiguration.Builder()
            .seed(123)  // 设置随机种子，确保结果可复现
            .weightInit(WeightInit.XAVIER)  // 使用Xavier权重初始化
            .updater(new Adam(0.001))  // 使用Adam优化器，学习率0.001
            .list()
            // 第一个卷积层：16个3x3的卷积核
            .layer(new ConvolutionLayer.Builder(3, 3)
                .nIn(图片通道数)
                .stride(1, 1)
                .nOut(16)
                .activation(Activation.RELU)
                .build())
            // 第一个最大池化层：2x2池化窗口
            .layer(new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                .kernelSize(2, 2)
                .stride(2, 2)
                .build())
            // 第二个卷积层：32个3x3的卷积核
            .layer(new ConvolutionLayer.Builder(3, 3)
                .stride(1, 1)
                .nOut(32)
                .activation(Activation.RELU)
                .build())
            // 第二个最大池化层：2x2池化窗口
            .layer(new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                .kernelSize(2, 2)
                .stride(2, 2)
                .build())
            // 全连接层：64个神经元
            .layer(new DenseLayer.Builder()
                .nOut(64)
                .activation(Activation.RELU)
                .build())
            // 输出层：使用softmax激活函数和负对数似然损失函数
            .layer(new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
                .nOut(病虫害种类数)
                .activation(Activation.SOFTMAX)
                .build())
            .setInputType(InputType.convolutional(图片高度, 图片宽度, 图片通道数))
            .build();

        模型 = new MultiLayerNetwork(配置);
        模型.init();
        模型.setListeners(new ScoreIterationListener(10));  // 每10次迭代输出一次损失值
    }

    /**
     * 训练模型
     * @param 训练数据 训练数据集
     * @param 训练轮数 训练的总轮数
     */
    public void train(DataSetIterator 训练数据, int 训练轮数) {
        for (int 当前轮数 = 0; 当前轮数 < 训练轮数; 当前轮数++) {
            System.out.println("第 " + (当前轮数 + 1) + " 轮训练 / 共 " + 训练轮数 + " 轮");
            模型.fit(训练数据);
            训练数据.reset();
        }
    }

    /**
     * 获取训练好的模型
     * @return 训练好的神经网络模型
     */
    public MultiLayerNetwork getModel() {
        return 模型;
    }
} 